import argparse


def args_parser():
    parser = argparse.ArgumentParser()
    # python main.py --algorithm=$1 --global_model=$2 --local_model=$3 --iid=$4 --dataset=$5 --results_save=$6

    # python main.py --algorithm=fed_avg --node_num=1 --global_model=mlp --local_model=mlp --E=5 --R=50
    # python main.py --algorithm=fed_mutual --node_num=5 --global_model=CNN --local_model=mlp --E=5 --R=50

    # Total
    parser.add_argument('--algorithm', type=str, default='fedavg',
                        help="Type of algorithms:{fml, fedavg,fedprox}")
    parser.add_argument('--device', type=str, default='cuda:0',
                        help="device: {cuda, cpu}")
    parser.add_argument('--num_users', type=int, default=5, help="number of users: K")
    parser.add_argument('--shard_per_user', type=int, default=6, help="classes per user")
    parser.add_argument('--R', type=int, default=200,
                        help="Number of rounds: R")
    parser.add_argument('--E', type=int, default=3,
                        help="Number of local epochs: E")
    parser.add_argument('--frac', type=float, default=1,  # 选择比例
                        help='the fraction of clients: C')
    parser.add_argument('--iid', type=int, default=1,  # 选择比例
                        help='set 1 for iid')
    parser.add_argument('--dataset', type=str, default='cifar100',
                        help="Type of algorithms:{mnist, cifar10,cifar100}")

    # Model
    parser.add_argument('--global_model', type=str, default='vgg',
                        help='Type of global customized_model: {mlp, LeNet5,CNN, ResNet18}')
    parser.add_argument('--local_model', type=str, default='vgg',
                        help='Type of local customized_model: {mlp, LeNet5,CNN, ResNet18}')
    parser.add_argument('--results_save', type=str, default='/', help='define fed results save folder')
    # Data
    parser.add_argument('--batchsize', type=int, default=256,
                        help="batchsize")

    # Optim
    parser.add_argument('--optimizer', type=str, default='sgd',
                        help="optimizer: {sgd, adam}")
    parser.add_argument('--lr', type=float, default=0.05,
                        help='learning rate')
    parser.add_argument('--lr_step', type=int, default=10,
                        help='learning rate decay step size')
    parser.add_argument('--stop_decay', type=int, default=100,
                        help='round when learning rate stop decay')
    parser.add_argument('--momentum', type=float, default=0.9,
                        help='SGD momentum')
    parser.add_argument('--alpha', type=float, default=0.5,
                        help='local ratio of data loss')
    parser.add_argument('--beta', type=float, default=0.5,
                        help='copy_meme ratio of data loss')
    parser.add_argument('--mu', type=float, default=0.01,
                        help="Number of local epochs: E")
    args = parser.parse_args()
    return args
